[JAX] Support for cuDNN-backed flex attention by vcherepanov-nv · Pull Request #2985 · NVIDIA/TransformerEngine

vcherepanov-nv · 2026-05-13T03:18:24Z

Description

Adds experimental JAX fused-attention score_mod support through cuDNN frontend SDPA graphs.

This introduces a score_mod(graph, score, tensors) callback path for fused_attn, plus optional score_mod_bprop(graph, dscore, tensors) support for backward. The Python side builds and serializes cuDNN frontend forward/backward graphs, caches graph metadata with stable callback keys, supports auxiliary tensor operands, and supports Python/NumPy scalar operands as cuDNN pass-by-value tensors. The C++ JAX extension deserializes and caches the graphs per device, then executes them through new forward/backward FFI handlers.

The Flax API now plumbs score_mod through DotProductAttention, MultiHeadAttention, and TransformerLayer. Packed QKV/KV layouts are unpacked to the separate BSHD layout when score modification is requested.

Users are responsible for supplying a mathematically correct score_mod_bprop for the corresponding score_mod; Transformer Engine wires the callback into the cuDNN graph but does not validate gradient semantics.

Current score_mod limitations:

Requires fused attention to be enabled.
Supports separate rank-4 BSHD_BSHD_BSHD Q/K/V tensors only.
Supports FP16/BF16 Q/K/V tensors.
Mutually exclusive with attention bias, masks, sequence descriptors, dropout, sliding-window attention, packed/ragged metadata, context parallelism, and non-vanilla softmax/softmax offset.
Requires matching cuDNN frontend Python package and C++ headers.

Fixes # (issue)
#2492

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Please list the changes introduced in this PR:

A new score_mod code path for the JAX FusedAttention backend
cuDNN frontend graph serialization and JAX FFI execution for score_mod forward/backward
Flax plumbing for DotProductAttention, MultiHeadAttention, and TransformerLayer
Tests

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

for more information, see https://pre-commit.ci

greptile-apps · 2026-05-13T03:29:19Z

Greptile Summary

This PR adds experimental cuDNN-frontend-backed flex attention (score_mod) to the JAX backend, including graph serialization/deserialization, a Python-level graph cache, new C++ FFI handlers for forward and backward, and Flax plumbing through DotProductAttention, MultiHeadAttention, and TransformerLayer.

Core path: fused_attn short-circuits to a new _fused_attn_score_mod custom_vjp primitive when score_mod is provided; the Python side builds and serializes cuDNN frontend graphs at trace time (cached by shape/dtype/config key), then passes the serialized bytes + UID maps as static FFI attributes to C++ handlers that deserialize and execute them.
C++ side: Two new FFI handlers deserialize graphs on demand into a process-lifetime unordered_map guarded by a mutex, with a thread-local cuDNN handle cache; the current double-checked locking leaves a window for redundant concurrent deserialization.
Flax plumbing: Packed and KV-packed layouts are transparently converted to separate BSHD tensors before the score_mod path; score_mod_tensors / score_mod_bprop_tensors are forwarded as call-time arguments to keep tensor operands in the JAX computation graph.

Confidence Score: 5/5

Safe to merge as an experimental feature; all flagged items are non-blocking quality improvements with no correctness impact.

The core forward and backward graph building, caching, FFI dispatch, and Flax plumbing are all structurally correct. Cache key stability, UID ordering, and pytree gradient structure are handled properly. The findings are race conditions that produce at worst redundant work (not wrong results) and a shutdown-order concern for the thread-local cuDNN handle that matches patterns already present elsewhere in the codebase.

transformer_engine/jax/csrc/extensions/attention.cpp (double-checked locking in GetScoreModGraph, thread-local handle destructor ordering) and transformer_engine/jax/cpp_extensions/flex_attention.py (Python-level cache lock).

Important Files Changed

Filename	Overview
transformer_engine/jax/cpp_extensions/flex_attention.py	New 967-line file implementing cuDNN frontend score_mod graph building, caching, and FFI dispatch; implements a stable cache-key scheme and separates tensor vs. scalar operands cleanly.
transformer_engine/jax/csrc/extensions/attention.cpp	Adds 251 lines for C++ cuDNN graph deserialization, thread-local handle cache, and two new FFI handlers (forward/backward); double-checked locking leaves redundant deserializations possible under thread contention.
transformer_engine/jax/attention.py	Adds custom_vjp wrapper for score_mod path with correct residual propagation, early-return before the deprecated sequence_descriptor path, and proper validation delegation.
transformer_engine/jax/flax/transformer.py	Plumbs score_mod/score_mod_bprop through DotProductAttention, MultiHeadAttention, and TransformerLayer; handles packed/kvpacked layout unpacking correctly before the score_mod path.
tests/jax/test_fused_attn_score_mod.py	New 671-line test suite covering causal masking, post-scale bias, softcap (forward/backward), and Flax layer integration, with reference implementations for correctness comparison.

Sequence Diagram

sequenceDiagram
    participant User
    participant fused_attn
    participant ScoreMod as "_fused_attn_score_mod"
    participant FlexPy as "flex_attention.py"
    participant FFI as "FFI/XLA"
    participant Cpp as "C++ Handler"
    participant Cache as "cuDNN Graph Cache"

    User->>fused_attn: "call with score_mod callback"
    fused_attn->>fused_attn: "validate_fused_attn_score_mod()"
    fused_attn->>FlexPy: "make_fused_attn_score_mod_config()"
    fused_attn->>ScoreMod: "custom_vjp forward"

    Note over ScoreMod,FlexPy: JAX Tracing Phase
    ScoreMod->>FlexPy: "fused_attn_score_mod_fwd()"
    FlexPy->>FlexPy: "check _score_mod_graph_cache"
    alt cache miss
        FlexPy->>FlexPy: "_build_score_mod_fwd_graph()"
        FlexPy->>FlexPy: "store in _score_mod_graph_cache"
    end
    FlexPy->>FFI: "ffi.ffi_call(serialized_graph, uids)"

    Note over FFI,Cache: XLA Execution Phase
    FFI->>Cpp: "FusedAttnScoreModForwardFFI(stream, q, k, v)"
    Cpp->>Cache: "GetScoreModGraph(stream, attrs)"
    alt C++ cache miss
        Cache->>Cache: "graph->deserialize(handle, data)"
        Cache->>Cache: "store shared_ptr in map"
    end
    Cpp->>Cpp: "graph->execute(handle, variant_pack)"
    Cpp-->>FFI: "output, stats, workspace"

    Note over ScoreMod,FlexPy: Backward pass
    ScoreMod->>FlexPy: "fused_attn_score_mod_bwd(qkv, o, dO, stats)"
    FlexPy->>FFI: "ffi.ffi_call(serialized_bwd_graph)"
    FFI->>Cpp: "FusedAttnScoreModBackwardFFI(...)"
    Cpp-->>FFI: "dq, dk, dv"

_{Reviews (9): Last reviewed commit: "Skip softcap score-mod test before SM90" | Re-trigger Greptile}

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

for more information, see https://pre-commit.ci

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

for more information, see https://pre-commit.ci

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

for more information, see https://pre-commit.ci

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

vcherepanov-nv added 2 commits May 12, 2026 05:15

Add JAX fused attention score_mod support

f967a26

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Stabilize score_mod callback cache keys

6b05328

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

vcherepanov-nv requested a review from cyanguwa May 13, 2026 03:18

vcherepanov-nv requested review from KshitijLakhani and jberchtold-nvidia as code owners May 13, 2026 03:18

vcherepanov-nv added the 2.16.0 label May 13, 2026

[pre-commit.ci] auto fixes from pre-commit.com hooks

1a96352

for more information, see https://pre-commit.ci

vcherepanov-nv mentioned this pull request May 13, 2026

[Draft]Support for score_mod and score_mod_bprop in cuDNN's sdpa #2767

Closed

13 tasks

greptile-apps Bot reviewed May 13, 2026

View reviewed changes

Add distributed JAX score mod attention test

3bf9e97

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

greptile-apps Bot reviewed May 13, 2026

View reviewed changes

Comment thread tests/jax/test_fused_attn.py Outdated

vcherepanov-nv and others added 2 commits May 15, 2026 03:35

Address JAX score_mod review items

29bbac7

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

c597af5

for more information, see https://pre-commit.ci

jberchtold-nvidia reviewed May 15, 2026

View reviewed changes

vcherepanov-nv and others added 2 commits May 18, 2026 23:54

Use serialized cuDNN graphs for score_mod attention

f8bd844

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

2c01c5e

for more information, see https://pre-commit.ci

jberchtold-nvidia reviewed May 19, 2026

View reviewed changes

Comment thread transformer_engine/jax/csrc/extensions/attention.cpp Outdated

Comment thread transformer_engine/jax/csrc/extensions/attention.cpp Outdated

vcherepanov-nv and others added 2 commits May 19, 2026 00:32

Rename score_mod graph cache helpers

ba6a1a7

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

deebf8e

for more information, see https://pre-commit.ci

KshitijLakhani requested changes May 20, 2026

View reviewed changes

Add Flax score_mod attention support

9a92dd2

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

jberchtold-nvidia reviewed May 21, 2026

View reviewed changes

Comment thread transformer_engine/jax/attention.py

Address JAX score_mod review feedback

2198a79

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

github-actions Bot added the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label May 21, 2026

vcherepanov-nv added 2 commits May 22, 2026 12:20

Add flex-attn tests to QA scripts

ffc8c79

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Skip softcap score-mod test before SM90

8856323

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[JAX] Support for cuDNN-backed flex attention#2985

[JAX] Support for cuDNN-backed flex attention#2985
vcherepanov-nv wants to merge 14 commits into
NVIDIA:mainfrom
vcherepanov-nv:cudnn-score-mod-jax

vcherepanov-nv commented May 13, 2026 •

edited

Loading

Uh oh!

greptile-apps Bot commented May 13, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

vcherepanov-nv commented May 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Checklist:

Uh oh!

greptile-apps Bot commented May 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Sequence Diagram

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

vcherepanov-nv commented May 13, 2026 •

edited

Loading

greptile-apps Bot commented May 13, 2026 •

edited

Loading